/**
 * 
 */
package codec;

import java.util.ArrayList;

import codec.huffman.JPEGHuffmanTable;
import util.BitSequenceV2;
import util.ColorTransformer;
import util.Log;
import util.Transformer;

import image.Image;
import image.header.JPEGHeader;

/**
 * @author nquangtrung
 *
 */
public class JPEGCodec {
	private Image sourceImage;
	private JPEGHuffmanTable huffmanTable;
	
	private JPEGHeader header;
	
	public static final String TAG = "JPEGCodec";

	public JPEGCodec(Image image, JPEGHeader header) {
		this.setHeader(header);
		this.setHuffmanTable(header.getHuffmanTable());
		this.setSourceImage(image);
	}
	
	private int[][][] prepareImage() {
		Image image = this.getSourceImage();
		
		ColorTransformer ct = new ColorTransformer();
		
		int width = image.getWidth();
		int height = image.getHeight();
		
		// Convert the image from RGB into YCbCr
		// Luminance - Chrominance Blue - Chrominance Red
		Log.v(TAG, "Coverting image from RGB to YCbCr");
		int[][][] ycc = ct.rgb2ycc(image);
		// Down sample the YCC image
		Log.v(TAG, "Down sample Image");
		if (this.getHeader().isDownsample())
			ycc = ct.yccdownsample(ycc, width, height);
		
		return ycc;
	}
	
	private int[][] prepareMCU(int[][][] mcu) {
		Transformer t = new Transformer();
		
		// With each block, do a DCT transform
		double[][][] g = new double[3][8][8];
		for (int i = 0; i < 3; i++) {
			g[i] = t.transform(mcu[i], 8, 8);
		}
		
		// And quantize it
		int[][][] q = t.quantize(g, 8, 8, this.getQuality());
		
		// Do zigzag order
		int[][] z = new int[3][64];
		for (int i = 0; i < 3; i++) {
			z[i] = t.zigzag(q[i], 8, 8);
		}
		return z;
	}
	
	protected ArrayList<Integer> makeWordList(int[][] z) {
		Transformer t = new Transformer();
		
		ArrayList<Integer> wordList = new ArrayList<Integer>();
		
		int[] oldDC = {0, 0, 0};
		
		for (int i = 0; i < 3; i++) {
			int[] component = z[i];
			int dc = component[0];
			int diff = dc - oldDC[i];
			wordList.add(diff);
			
			int lnz = t.lastNonZero(component, 64);
			int run = 0;
			for (int j = 1; j <= lnz; j++) {
				if (component[j] == 0) {
					run++;
					if (run == 16) {
						// ZRL
						int word = 0xF0;
						wordList.add(word);
					}
				} else {
					int acCoef = component[j];
					int size = JPEGHuffmanTable.size(acCoef);
					// (run, size) (addition)
					int word = run * 16 + size;
					wordList.add(word);
					run = 0;
				}
			}
			if (lnz < 63) {
				int word = 0x00; // eob
				wordList.add(word);
			}
		}
		
		return wordList;
	}
	
	public byte[] makeJPEGData() {
		Transformer t = new Transformer();
		int width = this.getSourceImage().getWidth();
		int height = this.getSourceImage().getHeight();
		
		// Make image ready for encode
		int[][][] ycc = this.prepareImage();
		
		int mcuCount = 0;
		int[] oldDC = {0, 0, 0};
		BitSequenceV2 jpegseq = new BitSequenceV2(width * height * 24);
		
		height = height - (height % 8);
		width = width - (width % 8);

		for (int y = height - 1; y >= 0; y -= 8) {
			for (int x = 0; x < width; x += 8) {
				
				Log.v(TAG, "=== Block #" + mcuCount + " ===========");
				mcuCount++;
				// Split into 8x8 block
				// start from top left
				Log.v(TAG, "Making block");
				int[][][] hexplet = new int[3][8][8];
				for (int dx = 0; dx < 8; dx++) {
					for (int dy = 0; dy < 8; dy++) {
						hexplet[0][dx][dy] = ycc[0][x+dx][y-dy];
						hexplet[1][dx][dy] = ycc[1][x+dx][y-dy];
						hexplet[2][dx][dy] = ycc[2][x+dx][y-dy];
					}
				}
				
				// Preparing the mcu
				int[][] z = this.prepareMCU(hexplet);
				
				Log.v(TAG, "Huffman encode: ");
				int[] index = new int[3];
				// Encode each component
				// Luminance
				int[] codeCount = {1, 1, 1};
				for (int i = 0; i < 3; i++) {
					Log.v(TAG, "==== Component #" + i + " =======");
					index[i] = 0;
					int[] quantized_y = z[i];
					int dc = quantized_y[0];
					int diff = dc - oldDC[i];
					oldDC[i] = dc;
					
					int size = JPEGHuffmanTable.size(diff);
					Log.v(TAG, "DC Size: " + size + " diff: " + diff);
					
					int[] cod = new int[width * height];
					int[] col = new int[width * height];
					
					index[i] += extendCode(cod, col, 0, size, i, 1, size, diff);
					
					int lnz = t.lastNonZero(quantized_y, 64);
					int run = 0;
					for (int j = 1; j <= lnz; j++) {
						if (quantized_y[j] == 0) {
							run++;
							if (run == 16) {
								// ZRL
								int word = 0xF0;
								index[i] += extendCode(cod, col, index[i], word, i, 0, 0, 0);
								
								run = 0;
								
								codeCount[i] += 15;
							}
						} else {
							int acCoef = quantized_y[j];
							size = JPEGHuffmanTable.size(acCoef);
							// (run, size) (addition)
							int word = run * 16 + size;
							Log.v(TAG, "Word: " + word + " run: " + run + " size: " + size + " ac: " + acCoef);
							
							index[i] += extendCode(cod, col, index[i], word, i, 0, size, acCoef);
							
							codeCount[i] += run + 1;
							run = 0;
						}
					}
					if (quantized_y[63] == 0) {
						int word = 0x00; // eob
						index[i] += extendCode(cod, col, index[i], word, i, 0, 0, 0);
						codeCount[i] += 1;
					}
					
					Log.debugArray(TAG, "Code:", cod, 8);
					Log.debugArray(TAG, "Code Length: ", col, 8);
					// End of coding procedure
					for (int j = 0; j < index[i]; j++) {
						jpegseq.stuff(cod[j], col[j]);
					}
				}
				boolean b = false;
				if (b) {
					System.out.printf(" (%3dx%3d)%2d-%2d-%2d ", x, height - y, codeCount[0], codeCount[1], codeCount[2]);
					if (mcuCount % 8 == 0) System.out.println();
				}
			}
		}
		byte[] jpegscan = jpegseq.makeByte();
		int bitcount = jpegseq.getCurrentBitOffset();
		// bit stuffing
		byte[] bitstuff = new byte[jpegscan.length * 10];
		for (int i = 0, index = 0; i < jpegscan.length; i++, index++) {
			bitstuff[index] = jpegscan[i];
			if (jpegscan[i] == (byte) 0xFF) {
				index++;
				bitstuff[index] = 0x00;
				bitcount += 8;
			}
		}
				
		int lnz = (bitcount - 1) / 8;
		byte[] finalscan = new byte[lnz + 1];
		for (int i = 0; i <= lnz; i++) {
			finalscan[i] = bitstuff[i];
		}
		// System.out.println(bitcount);
		Log.debugArray(TAG, "Final scan: ", finalscan, 8);
		
		return finalscan;
	}
	private int extendCode(int[] code, int[] codelength, int offset, int word, int component, int type, int size, int addtionalBit) {
		JPEGHuffmanTable tbl = this.getHuffmanTable();
		int table = 0;
		if (type == 0) {				// type = 0 => AC
			if (component == 0) {		// luminance
				table = JPEGHuffmanTable.LUMINANCE_AC;
			} else {					// chrominance
				table = JPEGHuffmanTable.CHROMINANCE_AC;
			}
		} else {						// else => DC
			if (component == 0) {		// luminance
				table = JPEGHuffmanTable.LUMINANCE_DC;
			} else {					// chrominance
				table = JPEGHuffmanTable.CHROMINANCE_DC;
			}
		}

		int c = tbl.findCode(word, table);
		int s = tbl.findSize(word, table);
		
		code[offset] = c;
		codelength[offset] = s;
		// additional bit
		if (word != 0xF0 && word != 0x00) {
			// if word is not zrl or eob
			code[offset + 1] = (addtionalBit>0)?addtionalBit:(addtionalBit - 1);
			codelength[offset + 1] = size;
			
			return 2;
		}
		return 1;
	}
	/**
	 * @return the quality
	 */
	public double getQuality() {
		return this.getHeader().getQuality();
	}
	/**
	 * @return the sourceImage
	 */
	public Image getSourceImage() {
		return sourceImage;
	}
	/**
	 * @param sourceImage the sourceImage to set
	 */
	public void setSourceImage(Image sourceImage) {
		this.sourceImage = sourceImage;
	}
	/**
	 * @return the huffmanTable
	 */
	public JPEGHuffmanTable getHuffmanTable() {
		return huffmanTable;
	}
	/**
	 * @param huffmanTable the huffmanTable to set
	 */
	public void setHuffmanTable(JPEGHuffmanTable huffmanTable) {
		this.huffmanTable = huffmanTable;
	}

	/**
	 * @return the header
	 */
	public JPEGHeader getHeader() {
		return header;
	}

	/**
	 * @param header the header to set
	 */
	public void setHeader(JPEGHeader header) {
		this.header = header;
	}
}